# 接收文件API， 文件分割，将文件保存到向量库中
from flask import Blueprint, request, jsonify,Response
import os,json
from .ollama_service import OllamaService

UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

file_rag_bp = Blueprint('rag', __name__)

ollama_service = OllamaService()

@file_rag_bp.route('/upload', methods=['POST'])
def upload():
    if 'document' not in request.files:
        return jsonify({'error': 'document mast be have !'}), 400

    document = request.files['document']
    if not document.filename.endswith('.pdf'):
        return jsonify({'error': 'Only PDF files are allowed!'}), 400

    # Save the image to a temporary location
    document_path = os.path.join(UPLOAD_FOLDER, document.filename)
    document.save(document_path)
    try:
        ollama_service.add_documents(document_path)
        return jsonify({'success': 'Document saved and processed successfully!'})
    except Exception as e:
        print(e)
        return jsonify({'error': str(e)}), 500

@file_rag_bp.route('/ask', methods=['POST'])
def ask():
    question = request.json.get('question')
    print(question)
    try:
        def generate_stream():
            # 修改为流式调用方式
            response_stream = ollama_service.qa_chain.stream({"question":question})
            for idx, chunk in enumerate(response_stream):
                result= {
                    "result":{
                        "output":{
                            "text": ''.join(chunk)
                        },
                        "metadata": {
                            "finishReason": "stream"
                        }
                    }
                }
                print(chunk, end="", flush=True)
                yield f"id: {idx}\ndata: {json.dumps(result, ensure_ascii=False)}\n\n"
            result= {
                "result":{
                    "output":{
                        "text": ''.join(chunk)
                    },
                    "metadata": {
                        "finishReason": "stop"
                    }
                }
            }
            yield f"id: {idx}\ndata: {json.dumps(result, ensure_ascii=False)}\n\n"

        return Response(generate_stream(),mimetype="text/event-stream")
    except Exception as e:
        return jsonify({'error': str(e)}), 500
@file_rag_bp.route('/delete/<path:filename>', methods=['GET'])
def delete_file(filename):
    os.remove(filename)
    return jsonify({'success': 'Document deleted successfully!'})